package grpc_proxy_middleware

import (
	"strings"

	"gin-gateway/core/dao"
	"github.com/pkg/errors"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
)

// 匹配接入方式 基于请求信息
func GrpcHeaderTransferMiddleware(serviceDetail *dao.ServiceDetail) func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
	return func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
		md, ok := metadata.FromIncomingContext(ss.Context())
		if !ok {
			return errors.New("miss metadata from context")
		}
		if len(serviceDetail.GRPCRule.HeaderTransfor) != 0 {
			for _, item := range strings.Split(serviceDetail.GRPCRule.HeaderTransfor, ",") {
				items := strings.Split(item, " ")
				if len(items) != 3 {
					continue
				}
				if items[0] == "add" || items[0] == "edit" {
					md.Set(items[1], items[2])
				}
				if items[0] == "del" {
					delete(md, items[1])
				}
			}

			// 不能使用 ss.SetHeader 方法 ！！！
			// 只能通过修改 ss.Context 的方式修改传给向代理 rpc 的 metadata, 而且必须用 metadata.NewIncomingContext 方式新建 context,
			ctx := metadata.NewIncomingContext(ss.Context(), md)
			ss = &wrappedSS{ServerStream: ss, Ctx: ctx}
		}
		if err := handler(srv, ss); err != nil {
			// log.Printf("RPC failed with error %v\n", err)
			return err
		}
		return nil
	}
}
